package nats

import (
	"context"
	"sync"

	"gitee.com/mkwchecking/sr-core/mq"
	"gitee.com/mkwchecking/sr-core/yal/yconv"
	"github.com/tx7do/kratos-transport/broker"
	"github.com/tx7do/kratos-transport/broker/nats"
)

var _ mq.IMsgQueue = (*NatsMq)(nil)

func init() {
	mq.Register(mq.Nats, NewNatsMq)
}

func NewNatsMq(conn mq.Conn) (mq.IMsgQueue, error) {

	broker := nats.NewBroker(
		broker.WithAddress(conn.Address),
		broker.WithCodec("json"),
	)

	if err := broker.Init(); err != nil {
		return nil, err
	}

	if err := broker.Connect(); err != nil {
		return nil, err
	}

	s := &NatsMq{
		broker:     broker,
		subscriber: sync.Map{},
	}
	return s, nil
}

type NatsMq struct {
	broker     broker.Broker
	subscriber sync.Map
}

// Publish implements mq.IMsgQueue.
func (n *NatsMq) Publish(ctx context.Context, msg mq.MQMessage) error {
	return n.broker.Publish(ctx, msg.Topic, msg)
}

// Subscribe implements mq.IMsgQueue.
func (n *NatsMq) Subscribe(ctx context.Context, topic string, handlers ...mq.IMsgHandler) error {
	subscriber, err := n.broker.Subscribe(topic, func(ctx context.Context, msg broker.Event) error {
		mqMsg, err := yconv.AnyToObjG[mq.MQMessage](msg.Message().Body)
		if err != nil {
			return err
		}

		for _, handler := range handlers {
			if err := handler.Handle(ctx, mqMsg); err != nil {
				return err
			}
		}
		return nil
	}, nil)

	if err != nil {
		return err
	}
	n.addScriber(topic, subscriber)
	return nil
}

// Unsubscribe implements mq.IMsgQueue.
func (n *NatsMq) Unsubscribe(ctx context.Context, topic string, handlers ...mq.IMsgHandler) error {
	subscriber, ok := n.getSubscriber(topic)
	if !ok {
		return nil
	}

	if err := subscriber.Unsubscribe(true); err != nil {
		return err
	}
	n.removeScriber(topic)
	return nil
}

func (n *NatsMq) addScriber(topic string, subscriber broker.Subscriber) {
	n.subscriber.Store(topic, subscriber)
}

func (n *NatsMq) removeScriber(topic string) {
	n.subscriber.Delete(topic)
}

func (n *NatsMq) getSubscriber(topic string) (broker.Subscriber, bool) {
	subscriber, ok := n.subscriber.Load(topic)
	if !ok {
		return nil, false
	}
	return subscriber.(broker.Subscriber), true
}
